from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())

loader = PyPDFLoader("../llama2.pdf")
pages = loader.load_and_split()

# print(pages[0].page_content)


text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=200,
    chunk_overlap=100,  # 思考：为什么要做overlap
    length_function=len,
    add_start_index=True,
)

# paragraphs = text_splitter.create_documents([pages[0].page_content])
# for para in paragraphs:
# print(para.page_content)
# print('-------')

texts = text_splitter.create_documents(
    [pages[2].page_content, pages[3].page_content])

# 灌库
embeddings = OpenAIEmbeddings()
db = Chroma.from_documents(texts, embeddings)

# 检索 top-1 结果
retriever = db.as_retriever(search_kwargs={"k": 1})

docs = retriever.get_relevant_documents("llama 2有多少参数？")

print(docs[0].page_content)
